import sys
import numpy as np
import enum
import json
from math import *

# 门矩阵编号类
class GateId(enum.Enum):
    i = 1
    x = 2
    y = 3
    z = 4
    h = 5
    s = 6
    sdg = 7
    t = 8
    tdg = 9
    rx = 10
    ry = 11
    rz = 12
    cx = 13
    cy = 14
    cz = 15
    swap = 16
    measure = 17

# 基本门矩阵类
class GateMatrix:
#     maxId = max([gid.value for gid in GateId.__members__.values()])
    gates = [] # 记录所有的门
    
    @staticmethod
    def init_i_matrix():
        GateMatrix.i = np.identity(2)
        GateMatrix.gates.append(GateMatrix.i)
        
    @staticmethod
    def init_x_matrix():
        GateMatrix.x = np.array([[0, 1], [1, 0]])
        GateMatrix.gates.append(GateMatrix.x)

    @staticmethod
    def init_y_matrix():
        GateMatrix.y = np.array([[0, 1j], [1j, 0]])
        GateMatrix.gates.append(GateMatrix.y)

    @staticmethod
    def init_z_matrix():
        GateMatrix.z = np.array([[1, 0], [0, -1]])
        GateMatrix.gates.append(GateMatrix.z)
    
    @staticmethod
    def init_h_matrix():
        GateMatrix.h = np.array([[1, 1], [1, -1]])
        GateMatrix.h = 1/sqrt(2) * GateMatrix.h
        GateMatrix.gates.append(GateMatrix.h)
        
    @staticmethod
    def init_s_matrix():
        GateMatrix.s = np.array([[1, 0], [0, 1j]])
        GateMatrix.gates.append(GateMatrix.s)
        
    @staticmethod
    def init_sdg_matrix():
        GateMatrix.sdg = np.array([[1, 0], [0, -1j]])
        GateMatrix.gates.append(GateMatrix.sdg)

    @staticmethod
    def init_t_matrix():
        GateMatrix.t = np.array([[1, 0], [0, e**(1j*pi/4)]])
        GateMatrix.gates.append(GateMatrix.t)

    @staticmethod
    def init_tdg_matrix():
        GateMatrix.tdg = np.array([[1, 0], [0, e**(-1j*pi/4)]])
        GateMatrix.gates.append(GateMatrix.tdg)

    @staticmethod
    def init_rx_matrix():
        GateMatrix.rx = np.zeros((2, 2))
        GateMatrix.gates.append(GateMatrix.rx)
    
    @staticmethod
    def init_ry_matrix():
        GateMatrix.ry = np.zeros((2, 2))
        GateMatrix.gates.append(GateMatrix.ry)

    @staticmethod
    def init_rz_matrix():
        GateMatrix.rz = np.zeros((2, 2))
        GateMatrix.gates.append(GateMatrix.rz)
        
    @staticmethod
    def update_r_matrix(gid, param):
        if (gid == GateId.rx.value):
            GateMatrix.rx = np.array([[cos(param/2), -1j*sin(param/2)], [-1j*sin(param/2), cos(param/2)]])
            GateMatrix.gates[gid] = GateMatrix.rx
        elif (gid == GateId.ry.value):
            GateMatrix.ry = np.array([[cos(param/2), -sin(param/2)], [sin(param/2), cos(param/2)]])
            GateMatrix.gates[gid] = GateMatrix.ry
        elif (gid == GateId.rz.value):
            GateMatrix.rz = np.array([[e**(-1j*param/2), 0], [0, e**(1j*param/2)]])
            GateMatrix.gates[gid] = GateMatrix.rz
        return
    
    @staticmethod
    def init_gate_matrix():
        # 初始化所有单比特输入无参数的门
        GateMatrix.init_i_matrix() # 0 留空
        GateMatrix.init_i_matrix()
        GateMatrix.init_x_matrix()
        GateMatrix.init_y_matrix()
        GateMatrix.init_z_matrix()
        GateMatrix.init_h_matrix()
        GateMatrix.init_s_matrix()
        GateMatrix.init_sdg_matrix()
        GateMatrix.init_t_matrix()
        GateMatrix.init_tdg_matrix()
        GateMatrix.init_rx_matrix()
        GateMatrix.init_ry_matrix()
        GateMatrix.init_rz_matrix()
#         print('===== init gate matrix =====')
#         for e in GateMatrix.gates:
#             print(e)
#         print('============================')

class QuantumCircuit:
    '''
    量子线路初始化
    nqubit: 线路的量子比特数
    '''
    def __init__(self, nqubit):
        # qubit(n)
        self.__nqubit = nqubit
        self.__meaqubit = np.zeros(self.__nqubit, dtype=np.int8) # 记录哪些qubit需要测量
        # state-vector(2^n)
        self.__slen = (1 << nqubit)
        self.__svector = np.zeros(self.__slen, dtype=np.complex128)
        self.__svector[0] = 1
        # operation-matrix(2^n x 2^n)
        self.__omatrix = np.zeros((self.__slen, self.__slen))
        # init gate matrics
        GateMatrix.init_gate_matrix()
        return
    
    def cx(self, ctrlid, targid):
        ctrlmask = (1 << (self.__nqubit - 1) >> ctrlid) # 控制位为1，其他都为0
        targmask = (1 << (self.__nqubit - 1) >> targid) # 受控位为1，其他都为0
        
        lowmask = (1 << self.__nqubit - ctrlid - 1) - 1
        highmask = ((1 << self.__nqubit - 1) - 1) ^ lowmask
        
        for idx in range(1 << self.__nqubit - 1): # 可以少遍历一半
            low = idx & lowmask
            high = (idx & highmask) << 1
            row1 = high | ctrlmask | low
            if (row1 & targmask == 0): # 受控位为0
                row2 = row1 ^ targmask
                self.__svector[[row1, row2]] = self.__svector[[row2, row1]]
        return

    def cy(self, ctrlid, targid):
        ctrlmask = (1 << (self.__nqubit - 1) >> ctrlid) # 控制位为1，其他都为0
        targmask = (1 << (self.__nqubit - 1) >> targid) # 受控位为1，其他都为0
        
        lowmask = (1 << self.__nqubit - ctrlid - 1) - 1
        highmask = ((1 << self.__nqubit - 1) - 1) ^ lowmask
        
        for idx in range(1 << self.__nqubit - 1): # 可以少遍历一半
            low = idx & lowmask
            high = (idx & highmask) << 1
            row1 = high | ctrlmask | low
            if (row1 & targmask == 0): # 受控位为0
                row2 = row1 ^ targmask
                a = self.__svector[row1]
                b = self.__svector[row2]
                self.__svector[row1] = -1j * b
                self.__svector[row2] = 1j * a
        return

    def cz(self, ctrlid, targid):
        ctrlmask = (1 << (self.__nqubit - 1) >> ctrlid) # 控制位为1，其他都为0
        targmask = (1 << (self.__nqubit - 1) >> targid) # 受控位为1，其他都为0
        
        lowmask = (1 << self.__nqubit - ctrlid - 1) - 1
        highmask = ((1 << self.__nqubit - 1) - 1) ^ lowmask
        
        for idx in range(1 << self.__nqubit - 1): # 可以少遍历一半
            low = idx & lowmask
            high = (idx & highmask) << 1
            row = high | ctrlmask | low
            if (row & targmask): # 控制位和受控位都为1，受控位相位翻转
                self.__svector[row] *= -1
        return
    
    def swap(self, qid1, qid2):
        ctrlmask = (1 << (self.__nqubit - 1) >> qid1) # 控制位为1，其他都为0
        targmask = (1 << (self.__nqubit - 1) >> qid2) # 受控位为1，其他都为0
        
        lowmask = (1 << self.__nqubit - qid1 - 1) - 1
        highmask = ((1 << self.__nqubit - 1) - 1) ^ lowmask
        
        for idx in range(1 << self.__nqubit - 1): # 可以少遍历一半
            low = idx & lowmask
            high = (idx & highmask) << 1
            row1 = high | ctrlmask | low
            if (row1 & targmask == 0): # |10>
                row2 = row1 ^ ctrlmask ^ targmask # |01>
                self.__svector[[row1, row2]] = self.__svector[[row2, row1]]
        return
    
    def id(self, qid):
        return
    
    def x(self, qid):
        self.__mtp(qid, GateId.x.value) # 直接计算操作矩阵
        self.__mmv()
        return
    
    def y(self, qid):
        self.__mtp(qid, GateId.y.value) # 直接计算操作矩阵
        self.__mmv()
        return
    
    def z(self, qid):
        self.__mtp(qid, GateId.z.value) # 直接计算操作矩阵
        self.__mmv()
        return
    
    def h(self, qid):
        self.__mtp(qid, GateId.h.value) # 直接计算操作矩阵
        self.__mmv()
        return
    
    def s(self, qid):
        self.__mtp(qid, GateId.s.value) # 直接计算操作矩阵
        self.__mmv()
        return
    
    def sdg(self, qid):
        self.__mtp(qid, GateId.sdg.value) # 直接计算操作矩阵
        self.__mmv()
        return
    
    def t(self, qid):
        self.__mtp(qid, GateId.t.value) # 直接计算操作矩阵
        self.__mmv()
        return
    
    def tdg(self, qid):
        self.__mtp(qid, GateId.tdg.value) # 直接计算操作矩阵
        self.__mmv()
        return
    
    def rx(self, qid, param):
        GateMatrix.update_r_matrix(GateId.rx.value, param) # 更新rx
        self.__mtp(qid, GateId.rx.value)
        self.__mmv()
        return
    
    def ry(self, qid, param):
        GateMatrix.update_r_matrix(GateId.ry.value, param) # 更新ry
        self.__mtp(qid, GateId.ry.value)
        self.__mmv()
        return
    
    def rz(self, qid, param):
        GateMatrix.update_r_matrix(GateId.rz.value, param) # 更新rz
        self.__mtp(qid, GateId.rz.value)
        self.__mmv()
        return
    
    def measure(self, mealist):
        for qid in mealist:
            self.__meaqubit[qid] = 1
        res = self.__measure()
        print(res)
        return

    '''
    矩阵张量积，给定一个量子门，构造对整个线路的操作矩阵
    作用于单输入的量子门
    '''
    def __mtp(self, qid, gid):
        # 1. 操作矩阵初始化
        if (qid == 0): # 作用在qubit0上的gate
            self.__omatrix = np.array(GateMatrix.gates[gid]) # 初始化成gates[gid]
        else:
            self.__omatrix = np.array(GateMatrix.i) # 初始化成ID
        # 2. 开始做张量积
        for i in range(1, self.__nqubit):
            if (i == qid):
                self.__omatrix = np.kron(self.__omatrix, GateMatrix.gates[gid])
            else:
                self.__omatrix = np.kron(self.__omatrix, GateMatrix.i)
        return

    '''操作矩阵与状态向量的乘积'''
    def __mmv(self):
        self.__svector = np.dot(self.__omatrix, self.__svector)
        return
    
    def __measure(self):
        # 结果字典
        res = {}
        
        # 将meaqubit数组转化成二进制串，再转为十进制数
        pattern = "".join([str(x) for x in self.__meaqubit])
        meamask = int(pattern, 2) # 需要测量的状态的mask
        
        # 和state vector的每个状态按位与
        for state in range(self.__slen):
            # 按位与得到下标
            key = meamask & state
            # 加入对应的集合
            res[key] = res.get(key, 0)
            res[key] += abs(self.__svector[state])**2
            
        # 可处理部分测量，将key转成二进制串，概率变百分制
        def to_state(num):
            biStr = format(num, '0'+str(self.__nqubit)+'b') # 整数转二进制字符串并补0
            state = ""
            for i in range(len(pattern)):
                if (pattern[i] == '1'):
                    state += biStr[i]
            return state
        
        prob = {}
        for k, v in res.items():
            prob[to_state(k)] = round(100 * v, 2)
        return prob

    def simulate(self, circuit):
        stop = False # 停止标记
        col = 0; # 列下标
        while ( not stop ):
            stop = True
            for qid in range(self.__nqubit): # 行下标
                # 1. 处理终止条件 #####
                if (circuit[qid] == None or col >= len(circuit[qid])): # 该qubit的量子门已经处理完毕
                    continue # 跳过，处理下一个qubit
                stop = False # 某一行还没处理完，不停止
                
                # 2. 调用模拟 #####
                gid = circuit[qid][col]['id']
                gtype = circuit[qid][col]['type']
                # 0) ID gate, do nothing
                if ( gtype == 0 ):
                    continue
                # 1) X, Y, Z, H, S, SDG, T, TDG
                elif ( gtype == 1 ):
                    self.__mtp(qid, gid) # 直接计算操作矩阵
                    self.__mmv()
                # 2) RX, RY, RZ
                elif ( gtype == 2 ):
                    gparam = circuit[qid][col]['params'] # 获取参数
                    GateMatrix.update_r_matrix(gid, gparam[0]) # 更新rx ry rz
                    self.__mtp(qid, gid)
                    self.__mmv()
                # 3) CX, CY, CZ, SWAP
                elif ( gtype == 3 ):
                    gparam = circuit[qid][col]['params'] # 获取参数
                    if ( gid == GateId.cx.value ):
                        self.cx(int(gparam[0]), int(gparam[1]))
                    elif ( gid == GateId.cy.value ):
                        self.cy(int(gparam[0]), int(gparam[1]))
                    elif ( gid == GateId.cz.value ):
                        self.cz(int(gparam[0]), int(gparam[1]))
                    elif ( gid == GateId.swap.value ):
                        self.swap(int(gparam[0]), int(gparam[1]))
                # 4) Measure
                elif ( gtype == 4 ):
                    self.__meaqubit[qid] = 1
                # ####################
            # 移动到后一列
            col += 1
        # 测量 #####
        return self.__measure() # 内部私有的measure
    
    def oracle():
        return

if __name__ == '__main__':
    # 1. 获取输入：线路抽象字符串转JSON #####
    cirJson = json.loads(sys.argv[1])
    # cirJson = json.loads(input())
    # 2. 调用模拟器 #####
    circ = QuantumCircuit(cirJson['nqubit'])
    res = circ.simulate(cirJson['circuit'])
    # 3. 返回概率向量JSON转字符串 #####
    print(json.dumps(res))

# 单输入：ID(0)，其他(1)
# 单输入，一个参数：rx ry rz(2)
# 双输入：swap cx cy cz(3)
# 测量：4